"""
# -*- coding: utf-8 -*-
# @Time    : 2023/5/27 19:15
# @Author  : 王摇摆
# @FileName: DIE.py
# @Software: PyCharm
# @Blog    ：https://blog.csdn.net/weixin_44943389?type=blog
"""
import numpy as np
import matplotlib.pyplot as plt


def butterfly_operation(x, N, k):
    """
    执行蝶形运算

    参数:
    x -- 输入信号
    N -- 信号长度
    k -- 蝶形运算的下标

    返回:
    输出信号
    """
    even_index = k
    odd_index = k + N // 2
    twiddle_factor = np.exp(-2j * np.pi * k / N)

    even_part = x[even_index]
    odd_part = x[odd_index] * twiddle_factor

    return even_part + odd_part, even_part - odd_part


def draw_butterfly(N):
    """
    绘制蝶形运算图

    参数:
    N -- 信号长度
    """
    num_stages = int(np.log2(N))
    num_butterflies = N // 2

    fig, axs = plt.subplots(num_stages, num_butterflies, figsize=(10, 6))
    fig.subplots_adjust(hspace=0.6)

    for stage in range(num_stages):
        for butterfly in range(num_butterflies):
            ax = axs[stage, butterfly]
            ax.cla()  # 清空当前子图

            ax.set_xlim([-1.5, 1.5])
            ax.set_ylim([-1.5, 1.5])

            # 绘制输入信号点
            ax.plot(0, 0, 'bo')

            # 绘制输出信号点
            ax.plot(1, 0, 'ro')
            ax.plot(1, 1, 'ro')

            # 计算蝶形运算结果
            x = np.zeros(N, dtype=np.complex128)
            x[butterfly], x[butterfly + N // 2] = 1, 1  # 输入信号设置为1
            y1, y2 = butterfly_operation(x, N, butterfly)

            # 绘制连接线
            ax.plot([0, 1], [0, 0], 'b-')  # 输入信号连接线
            ax.plot([1, 1], [0, 1], 'r-')  # 输出信号连接线1
            ax.plot([1, 1], [0, -1], 'r-')  # 输出信号连接线2

            # 绘制连接线上的箭头
            ax.arrow(0.5, 0, 0, -0.5, head_width=0.1, head_length=0.1, fc='b', ec='b')
            ax.arrow(1, 0.5, 0, -0.4, head_width=0.1, head_length=0.1, fc='r', ec='r')
            ax.arrow(1, -0.5, 0, 0.4, head_width=0.1, head_length=0.1, fc='r', ec='r')

            # 显示蝶形运算结果
            ax.text(1.1, 0.1, f'{y1:.2f}', color='r')
            ax.text(1.1, -1.1, f'{y2:.2f}', color='r')

            ax.set_title(f'Stage {stage + 1}, Butterfly {butterfly + 1}')

    plt.suptitle('Butterfly Operation')
    plt.show()


# 示例用法
N = 8  # 信号长度
draw_butterfly(N)
